package com.s3.demo.utils;

import com.amazonaws.ReadLimitInfo;
import com.amazonaws.SdkClientException;
import com.amazonaws.auth.AbstractAWSSigner;
import com.amazonaws.internal.SdkDigestInputStream;
import com.amazonaws.util.BinaryUtils;
import com.amazonaws.util.SdkHttpUtils;
import org.springframework.util.StringUtils;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.security.DigestInputStream;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.*;

import static com.amazonaws.util.StringUtils.UTF8;

/**
 * @author lzz
 * @create 2021-09-06 9:12
 */
public class AWSV4Auth {
    public static final String LINE_SEPARATOR = "\n";
    private static final List<String> listOfHeadersToIgnoreInLowerCase = Arrays.asList("connection", "x-amzn-trace-id");

    private static final String CONTENT_SHA_256 = "STREAMING-AWS4-HMAC-SHA256-PAYLOAD";
    private static final String UNSIGNED_PAYLOAD = "UNSIGNED-PAYLOAD";
    private AWSV4Auth() {
    }

    public static class Builder {

        private String accessKeyID;
        private String secretAccessKey;
        private String regionName;
        private String serviceName;
        private String httpMethodName;
        private String canonicalURI;
        private TreeMap<String, List<String>> queryParametes;
        private TreeMap<String, String> awsHeaders;
        private String payload;
        private String xAmzDate;
        private String currentDate;
        private boolean debug = false;

        public Builder(String accessKeyID, String secretAccessKey) {
            this.accessKeyID = accessKeyID;
            this.secretAccessKey = secretAccessKey;
        }

        public Builder regionName(String regionName) {
            this.regionName = regionName;
            return this;
        }

        public Builder serviceName(String serviceName) {
            this.serviceName = serviceName;
            return this;
        }

        public Builder httpMethodName(String httpMethodName) {
            this.httpMethodName = httpMethodName;
            return this;
        }

        public Builder canonicalURI(String canonicalURI) {
            this.canonicalURI = canonicalURI;
            return this;
        }

        public Builder queryParametes(TreeMap<String, List<String>> queryParametes) {
            this.queryParametes = queryParametes;
            return this;
        }

        public Builder awsHeaders(TreeMap<String, String> awsHeaders) {
            this.awsHeaders = awsHeaders;
            return this;
        }
        public Builder xAmzDate(String xAmzDate) {
            this.xAmzDate = xAmzDate;
            return this;
        }
        public Builder currentDate(String currentDate) {
            this.currentDate = currentDate;
            return this;
        }

        public Builder payload(String payload) {
            this.payload = payload;
            return this;
        }

        public Builder debug() {
            this.debug = true;
            return this;
        }

        public AWSV4Auth build() {
            return new AWSV4Auth(this);
        }
    }

    private String accessKeyID;
    private String secretAccessKey;
    private String regionName;
    private String serviceName;
    private String httpMethodName;
    private String canonicalURI;
    private TreeMap<String, List<String>> queryParametes;
    private TreeMap<String, String> awsHeaders;
    private String payload;
    private boolean debug = false;

    /* Other variables */
    private final String HMACAlgorithm = "AWS4-HMAC-SHA256";
    private final String aws4Request = "aws4_request";
    private String strSignedHeader;
    private String xAmzDate;
    private String currentDate;

    private AWSV4Auth(Builder builder) {
        accessKeyID = builder.accessKeyID;
        secretAccessKey = builder.secretAccessKey;
        regionName = builder.regionName;
        serviceName = builder.serviceName;
        httpMethodName = builder.httpMethodName;
        canonicalURI = builder.canonicalURI;
        queryParametes = builder.queryParametes;
        awsHeaders = builder.awsHeaders;
        payload = builder.payload;
        debug = builder.debug;

        /* Get current timestamp value.(UTC) */
        xAmzDate = getTimeStamp();
        currentDate = getDate();
        if (!StringUtils.isEmpty(builder.xAmzDate)) {
            xAmzDate = builder.xAmzDate;
        }

        if (!StringUtils.isEmpty(builder.currentDate)) {
            currentDate = builder.currentDate;
        }
    }

    /**
     * 任务 1：针对签名版本 4 创建规范请求
     *
     * @return
     */
    private String prepareCanonicalRequest(String contentHash) {
        StringBuilder canonicalURL = new StringBuilder("");

        /* Step 1.1 以HTTP方法(GET, PUT, POST, etc.)开头, 然后换行. */
        canonicalURL
                .append(httpMethodName)
                .append(LINE_SEPARATOR)
                /* Step 1.2 添加URI参数，换行. */
                .append(canonicalURI == null || canonicalURI.trim().isEmpty() ? "/" : canonicalURI)
                .append(LINE_SEPARATOR)
                /* Step 1.3 添加查询参数，换行. */
                .append(getCanonicalizedQueryString(queryParametes))
                .append(LINE_SEPARATOR)
                /* Step 1.4 添加headers, 每个header都需要换行. */
                /* Step 1.5 添加签名的headers并换行. */
                .append(getCanonicalizedHeaderString(awsHeaders))
                .append(LINE_SEPARATOR)
                // 添加 头 以;分割
                .append(getSignedHeadersString(awsHeaders))
                .append(LINE_SEPARATOR)
                .append(contentHash);

        if (debug) {
            System.out.println("##Canonical Request:\n" + canonicalURL.toString());
        }

        return canonicalURL.toString();
    }

    private String calculateContentHash() {
        if ("PUT".equals(httpMethodName)) {
            return CONTENT_SHA_256;
        }
        InputStream payloadStream =  new ByteArrayInputStream(new byte[0]);;
        return BinaryUtils.toHex(doStreamhash(payloadStream));
    }

    /**
     *  stream hash
     * @param input
     * @return
     * @throws SdkClientException
     */
    protected byte[] doStreamhash(InputStream input) throws SdkClientException {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            @SuppressWarnings("resource")
            DigestInputStream digestInputStream = new SdkDigestInputStream(input, md);
            byte[] buffer = new byte[1024];
            while (digestInputStream.read(buffer) > -1) ;
            return digestInputStream.getMessageDigest().digest();
        } catch (Exception e) {
            throw new SdkClientException(
                    "Unable to compute hash while signing request: "
                            + e.getMessage(), e);
        }
    }

    /**
     * 任务 2：创建签名版本 4 的待签字符串
     *
     * @param canonicalURL
     * @return
     */
    private String prepareStringToSign(String canonicalURL) throws NoSuchAlgorithmException {
        /* Step 2.1 以算法名称开头，并换行. */
        StringBuilder string2signBuilder = new StringBuilder(HMACAlgorithm);
        string2signBuilder.append(LINE_SEPARATOR)
                /* Step 2.2 添加日期，并换行. */
                .append(xAmzDate)
                .append(LINE_SEPARATOR)
                //scope currentDate/regionName/ServiceName/aws4_request
                /* Step 2.3 添加认证范围，并换行. */
                .append(currentDate+"/"+regionName+"/"+serviceName+"/"+aws4Request)
                .append(LINE_SEPARATOR)
                /* Step 2.4 添加任务1返回的规范URL哈希处理结果，然后换行. */
                .append(BinaryUtils.toHex(doHash(canonicalURL)));
        return string2signBuilder.toString();
    }

    /**
     * 任务 3：为 AWS Signature 版本 4 计算签名
     *
     * @param stringToSign
     * @return
     */
    private String calculateSignature(String stringToSign) {
        try {
            /* Step 3.1 生成签名的key */
            byte[] signatureKey = getSignatureKey(secretAccessKey, currentDate, regionName, serviceName);

            /* Step 3.2 计算签名. */
            byte[] signature = HmacSHA256(signatureKey, stringToSign);

            /* Step 3.2.1 对签名编码处理 */
            String strHexSignature = bytesToHex(signature);
            return strHexSignature;
        } catch (Exception ex) {
            ex.printStackTrace();
        }
        return null;
    }

    /**
     * 获取签名
     * @return
     * @throws NoSuchAlgorithmException
     */
    public String getSignature() throws NoSuchAlgorithmException {
        awsHeaders.put("x-amz-date", xAmzDate);

        /* 执行任务 1: 创建aws v4签名的规范请求字符串. */
        String canonicalURL = prepareCanonicalRequest(calculateContentHash());

        /* 执行任务 2: 创建用来认证的字符串 4. */
        String stringToSign = prepareStringToSign(canonicalURL);

        /* 执行任务 3: 计算签名. */
        return calculateSignature(stringToSign);
    }

    /**
     * 获取预签名
     * @return
     * @throws NoSuchAlgorithmException
     */
    public String getPreSignature() throws NoSuchAlgorithmException {

        /* 执行任务 1: 创建aws v4签名的规范请求字符串. */
        String canonicalURL = prepareCanonicalRequest(calculateContentHashPresign());

        /* 执行任务 2: 创建用来认证的字符串 4. */
        String stringToSign = prepareStringToSign(canonicalURL);

        /* 执行任务 3: 计算签名. */
        return calculateSignature(stringToSign);
    }

    public String calculateContentHashPresign() {
        return "UNSIGNED-PAYLOAD";
    }
    /**
     *任务 4：将签名信息添加到请求并返回headers
     *
     * @return
     */
    public Map<String, String> getHeaders() throws NoSuchAlgorithmException {
        awsHeaders.put("x-amz-date", xAmzDate);

        /* 执行任务 1: 创建aws v4签名的规范请求字符串. */
        String canonicalURL = prepareCanonicalRequest(calculateContentHash());

        /* 执行任务 2: 创建用来认证的字符串 4. */
        String stringToSign = prepareStringToSign(canonicalURL);

        /* 执行任务 3: 计算签名. */
        String signature = calculateSignature(stringToSign);

        if (signature != null) {
            Map<String, String> header = new HashMap<String, String>(0);
            header.put("x-amz-date", xAmzDate);
            strSignedHeader = getSignedHeadersString(awsHeaders);
            String authorizationString = buildAuthorizationString(signature);
            header.put("Authorization", authorizationString);
            if (debug) {
                System.out.println("##Signature:\n" + signature);
                System.out.println("##Header:");
                for (Map.Entry<String, String> entrySet : header.entrySet()) {
                    System.out.println(entrySet.getKey() + " = " + entrySet.getValue());
                }
                System.out.println("================================");
            }
            return header;
        } else {
            if (debug) {
                System.out.println("##Signature:\n" + signature);
            }
            return null;
        }
    }

    /**
     * 连接前几步处理的字符串生成Authorization header值.
     *
     * @param strSignature
     * @return
     */
    private String buildAuthorizationString(String strSignature) {
        return HMACAlgorithm + " "
                + "Credential=" + accessKeyID + "/" + getDate() + "/" + regionName + "/" + serviceName + "/" + aws4Request + ", "
                + "SignedHeaders=" + strSignedHeader + ", "
                + "Signature=" + strSignature;
    }

    /**
     * 将字符串16进制化.
     *
     * @param data
     * @return
     */
    private String generateHex(String data) {
        MessageDigest messageDigest;
        try {
            messageDigest = MessageDigest.getInstance("SHA-256");
            messageDigest.update(data.getBytes("UTF-8"));
            byte[] digest = messageDigest.digest();
            return String.format("%064x", new java.math.BigInteger(1, digest));
        } catch (NoSuchAlgorithmException | UnsupportedEncodingException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 以给定的key应用HmacSHA256算法处理数据.
     *
     * @param data
     * @param key
     * @return
     * @throws Exception
     * @reference:
     * http://docs.aws.amazon.com/general/latest/gr/signature-v4-examples.html#signature-v4-examples-java
     */
    private byte[] HmacSHA256(byte[] key, String data) throws Exception {
        String algorithm = "HmacSHA256";
        Mac mac = Mac.getInstance(algorithm);
        mac.init(new SecretKeySpec(key, algorithm));
        return mac.doFinal(data.getBytes("UTF8"));
    }

    /**
     * 生成AWS 签名
     *
     * @param key
     * @param date
     * @param regionName
     * @param serviceName
     * @return
     * @throws Exception
     * @reference
     * http://docs.aws.amazon.com/general/latest/gr/signature-v4-examples.html#signature-v4-examples-java
     */
    private byte[] getSignatureKey(String key, String date, String regionName, String serviceName) throws Exception {
        byte[] kSecret = ("AWS4" + key).getBytes("UTF8");
        byte[] kDate = HmacSHA256(kSecret, date);
        byte[] kRegion = HmacSHA256(kDate, regionName);
        byte[] kService = HmacSHA256(kRegion, serviceName);
        return HmacSHA256(kService, aws4Request);
    }

    /**
     * 头信息 key 以;分割
     * @param headers
     * @return
     */
    private String getSignedHeadersString(Map<String,String> headers) {
        final List<String> sortedHeaders = new ArrayList<String>(headers.keySet());
        Collections.sort(sortedHeaders, String.CASE_INSENSITIVE_ORDER);

        StringBuilder buffer = new StringBuilder();
        for (String header : sortedHeaders) {
            if (shouldExcludeHeaderFromSigning(header)) {
                continue;
            }
            if (buffer.length() > 0) {
                buffer.append(";");
            }
            buffer.append(com.amazonaws.util.StringUtils.lowerCase(header));
        }

        return buffer.toString();
    }

    /**
     * 标准头生成
     * @param requestHeaders
     * @return
     */
    private String getCanonicalizedHeaderString(Map<String, String> requestHeaders) {
        final List<String> sortedHeaders = new ArrayList<String>(requestHeaders
                .keySet());
        Collections.sort(sortedHeaders, String.CASE_INSENSITIVE_ORDER);

        StringBuilder buffer = new StringBuilder();
        for (String header : sortedHeaders) {
            if (shouldExcludeHeaderFromSigning(header)) {
                continue;
            }
            String key = com.amazonaws.util.StringUtils.lowerCase(header);
            String value = requestHeaders.get(header);

            com.amazonaws.util.StringUtils.appendCompactedString(buffer, key);
            buffer.append(":");
            if (value != null) {
                com.amazonaws.util.StringUtils.appendCompactedString(buffer, value);
            }

            buffer.append("\n");
        }

        return buffer.toString();
    }

    protected boolean shouldExcludeHeaderFromSigning(String header) {
        return listOfHeadersToIgnoreInLowerCase.contains(header.toLowerCase());
    }
    /**
     *  生成标准的查询参数字符串
     * @param parameters
     * @return
     */
    private String getCanonicalizedQueryString(Map<String, List<String>> parameters) {

        final SortedMap<String, List<String>> sorted = new TreeMap<String, List<String>>();

        /**
         * Signing protocol expects the param values also to be sorted after url
         * encoding in addition to sorted parameter names.
         */
        for (Map.Entry<String, List<String>> entry : parameters.entrySet()) {
            final String encodedParamName = SdkHttpUtils.urlEncode(
                    entry.getKey(), false);
            final List<String> paramValues = entry.getValue();
            final List<String> encodedValues = new ArrayList<String>(
                    paramValues.size());
            for (String value : paramValues) {
                encodedValues.add(SdkHttpUtils.urlEncode(value, false));
            }
            Collections.sort(encodedValues);
            sorted.put(encodedParamName, encodedValues);

        }
        final StringBuilder result = new StringBuilder();
        for(Map.Entry<String, List<String>> entry : sorted.entrySet()) {
            for(String value : entry.getValue()) {
                if (result.length() > 0) {
                    result.append("&");
                }
                result.append(entry.getKey())
                        .append("=")
                        .append(value);
            }
        }

        return result.toString();
    }

    final protected static char[] hexArray = "0123456789ABCDEF".toCharArray();

    /**
     * 将字节数组转换为16进制字符串
     *
     * @param bytes
     * @return
     */
    private String bytesToHex(byte[] bytes) {
        char[] hexChars = new char[bytes.length * 2];
        for (int j = 0; j < bytes.length; j++) {
            int v = bytes[j] & 0xFF;
            hexChars[j * 2] = hexArray[v >>> 4];
            hexChars[j * 2 + 1] = hexArray[v & 0x0F];
        }
        return new String(hexChars).toLowerCase();
    }

    /**
     * 获取hash
     * @param text
     * @return
     * @throws NoSuchAlgorithmException
     */
    private byte[] doHash(String text) throws NoSuchAlgorithmException {
        MessageDigest md = MessageDigest.getInstance("SHA-256");
        md.update(text.getBytes(UTF8));
        return md.digest();
    }

    /**
     * 获取yyyyMMdd'T'HHmmss'Z'格式的当前时间
     *
     * @return
     */
    private String getTimeStamp() {
        DateFormat dateFormat = new SimpleDateFormat("yyyyMMdd'T'HHmmss'Z'");
        //server timezone
        dateFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
        return dateFormat.format(new Date());
    }

    /**
     * 获取yyyyMMdd格式的当前日期
     *
     * @return
     */
    private String getDate() {
        DateFormat dateFormat = new SimpleDateFormat("yyyyMMdd");
        //server timezone
        dateFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
        return dateFormat.format(new Date());
    }

    /**
     * UTF-8编码
     * @param param
     * @return
     */
    private String encodeParameter(String param){
        try {
            return URLEncoder.encode(param, "UTF-8");
        } catch (Exception e) {
            return URLEncoder.encode(param);
        }
    }
}
